# import pandas as pd
# import pickle
#
#
# # 使用pandas读取pkl文件
# def read_pkl_with_pandas(file_path):
#     try:
#         data = pd.read_pickle(file_path)
#         print("使用pandas读取的pkl文件内容：")
#         print(data.head())  # 打印前5行数据
#
#         # 输出数据形状
#         if hasattr(data, 'shape'):
#             print("\n数据形状：", data.shape)
#     except Exception as e:
#         print(f"使用pandas读取pkl文件时出错：{e}")
#
#
# # 使用pickle读取pkl文件
# def read_pkl_with_pickle(file_path):
#     try:
#         with open(file_path, 'rb') as file:
#             data = pickle.load(file)
#         print("\n使用pickle读取的pkl文件内容：")
#         print(data)
#
#         # 输出数据形状
#         if hasattr(data, 'shape'):
#             print("\n数据形状：", data.shape)
#     except Exception as e:
#         print(f"使用pickle读取pkl文件时出错：{e}")
#
#
# # 使用示例
# if __name__ == "__main__":
#     # 替换为你的 .pkl 文件路径
#     # file_path  = "E:\datasets\CASIA-B_HRNet\CASIA-B_HRNet\\001\\bg-01\\000\\000.pkl"
#     file_path = "E:\datasets\CASIA-B\CASIA-B_sil-pkl\\012\\nm-02\\144\\144.pkl"
#     # file_path = "E:\datasets\\3.basketball\\3.d(3)p(2)s(3)\\videos\Basketball-P-silpkl\\001\p-02\\000\\000.pkl"
#     # file_path = "E:\datasets\CAISA-B-png_pkl\\001\\bg-01\\000\\000.pkl"
#
#     # file_path = "E:\datasets\\2D_sil-php\\0000\camid0_videoid2\seq0\seq0.pkl"
#     # file_path = "E:\datasets\AAAA\\2D_Poses-php-pkl\\0000\camid0_videoid2\seq0\seq0.pkl"
#
#     # file_path = "E:\datasets\\2.Gait3D\\2D_Poses-pkl\\0000\camid0_videoid2\seq0\\seq0.pkl"
#     # file_path = "E:\datasets\\2.Gait3D\\2D_Silhouettes-pkl\\0000\camid0_videoid2\seq0\\seq0.pkl"
#
#     # 使用pandas读取
#     # read_pkl_with_pandas(file_path)
#
#     # 使用pickle读取
#     read_pkl_with_pickle(file_path)


import torch
import torch.nn as nn
from graphviz import Digraph


class CMGF(nn.Module):
    def __init__(self, cfgs=None, is_training=True):
        super().__init__()
        # 模型结构定义
        self.set_block1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True)
        )
        self.set_block2 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)
        )
        self.set_pool0 = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(128, 10)

    def forward(self, x):
        x = self.set_block1(x)
        x = self.set_block2(x)
        x = self.set_pool0(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

    def visualize_architecture(self, filename='cmgf_model', view=True):
        """生成模型架构图"""
        # 创建有向图对象
        dot = Digraph(comment='CMGF Model Architecture',
                      node_attr={'fontname': 'Helvetica,Arial,sans-serif', 'style': 'filled'},
                      edge_attr={'fontname': 'Helvetica,Arial,sans-serif'})
        dot.attr(rankdir='TB', size='8,8', dpi='300', bgcolor='transparent')

        # 设置节点样式
        dot.attr('node', shape='box', color='lightblue', style='filled,rounded')

        # 添加节点
        dot.node('input', 'Input\n(3x224x224)', fillcolor='white')
        dot.node('block1', 'Set Block 1\n[Conv2d, BN, ReLU]', fillcolor='lightblue')
        dot.node('block2', 'Set Block 2\n[Conv2d, BN, ReLU, MaxPool]', fillcolor='lightblue')
        dot.node('pool', 'Global Average Pooling', fillcolor='lightgreen')
        dot.node('fc', 'Fully Connected\n(128 -> 10)', fillcolor='lightyellow')
        dot.node('output', 'Output\n(10 classes)', fillcolor='white')

        # 添加边
        dot.edge('input', 'block1', label='64x224x224')
        dot.edge('block1', 'block2', label='128x112x112')
        dot.edge('block2', 'pool', label='128x1x1')
        dot.edge('pool', 'fc', label='128')
        dot.edge('fc', 'output')

        # 渲染图形
        try:
            dot.render(filename, view=view, cleanup=True, format='pdf')
            print(f"模型架构图已保存为: {filename}.pdf")
        except Exception as e:
            print(f"生成架构图时出错: {e}")
            print("请确保已安装graphviz软件和Python包")
            # 作为备选，返回DOT格式的源代码
            return dot.source


# 使用示例
if __name__ == "__main__":
    # 创建模型实例
    model = CMGF()

    # 生成架构图
    model.visualize_architecture()

    # 打印模型概要
    print(model)